kreuzberg 3.0.0__py3-none-any.whl → 3.0.1__py3-none-any.whl

kreuzberg/_extractors/_base.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from kreuzberg import ExtractionResult
+    from kreuzberg._types import ExtractionConfig
+
+
+class Extractor(ABC):
+    """Abstract base class for file content extraction.
+
+    This class provides the interface for different types of content extractors.
+    Subclasses are expected to implement the methods for extracting content
+    either asynchronously or synchronously and determining the supported MIME types.
+
+    Attributes:
+        SUPPORTED_MIME_TYPES: The set of supported mime types - all none abstract extractors must implement this.
+
+    Args:
+        mime_type: The MIME type that this extractor handles (e.g., "application/pdf").
+        config: Configuration options for the extraction process.
+    """
+
+    __slots__ = ("config", "mime_type")
+
+    SUPPORTED_MIME_TYPES: ClassVar[set[str]]
+
+    def __init__(self, mime_type: str, config: ExtractionConfig) -> None:
+        self.mime_type = mime_type
+        self.config = config
+
+    @abstractmethod
+    async def extract_bytes_async(self, content: bytes) -> ExtractionResult:
+        """Asynchronously extract content from a byte stream.
+
+        Args:
+            content: The byte content to extract.
+
+        Returns:
+            ExtractionResult: The extracted content along with metadata about the extraction.
+        """
+
+    @abstractmethod
+    async def extract_path_async(self, path: Path) -> ExtractionResult:
+        """Asynchronously extract content from a file located at the specified path.
+
+        Args:
+            path: The path to the file to process.
+
+        Returns:
+            ExtractionResult: The extracted content along with metadata about the extraction.
+        """
+
+    @abstractmethod
+    def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
+        """Synchronously extract content from a byte stream.
+
+        Args:
+            content: The byte content to extract.
+
+        Returns:
+            ExtractionResult: The extracted content along with metadata about the extraction.
+        """
+
+    @abstractmethod
+    def extract_path_sync(self, path: Path) -> ExtractionResult:
+        """Synchronously extract content from a file located at the specified path.
+
+        Args:
+            path: The path to the file to process.
+
+        Returns:
+            ExtractionResult: The extracted content along with metadata about the extraction.
+        """
+
+    @classmethod
+    def supports_mimetype(cls, mime_type: str) -> bool:
+        """Verify whether the extractor supports the given MIME type.
+
+        Args:
+            mime_type: The MIME type to check (e.g., "application/pdf").
+
+        Returns:
+            bool: True if the MIME type is supported, False otherwise.
+        """
+        return mime_type in cls.SUPPORTED_MIME_TYPES or any(
+            mime_type.startswith(supported_type) for supported_type in cls.SUPPORTED_MIME_TYPES
+        )

kreuzberg/_extractors/_html.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+import html_to_markdown
+from anyio import Path as AsyncPath
+
+from kreuzberg._extractors._base import Extractor
+from kreuzberg._mime_types import HTML_MIME_TYPE, MARKDOWN_MIME_TYPE
+from kreuzberg._types import ExtractionResult
+from kreuzberg._utils._string import normalize_spaces, safe_decode
+from kreuzberg._utils._sync import run_sync
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+class HTMLExtractor(Extractor):
+    SUPPORTED_MIME_TYPES: ClassVar[set[str]] = {HTML_MIME_TYPE}
+
+    async def extract_bytes_async(self, content: bytes) -> ExtractionResult:
+        return await run_sync(self.extract_bytes_sync, content)
+
+    async def extract_path_async(self, path: Path) -> ExtractionResult:
+        content = await AsyncPath(path).read_bytes()
+        return await run_sync(self.extract_bytes_sync, content)
+
+    def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
+        result = html_to_markdown.convert_to_markdown(safe_decode(content))
+        return ExtractionResult(content=normalize_spaces(result), mime_type=MARKDOWN_MIME_TYPE, metadata={}, chunks=[])
+
+    def extract_path_sync(self, path: Path) -> ExtractionResult:
+        content = path.read_bytes()
+        return self.extract_bytes_sync(content)

kreuzberg/_extractors/_image.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+import anyio
+from anyio import Path as AsyncPath
+
+from kreuzberg._extractors._base import Extractor
+from kreuzberg._mime_types import IMAGE_MIME_TYPES
+from kreuzberg._ocr import get_ocr_backend
+from kreuzberg._utils._tmp import create_temp_file
+from kreuzberg.exceptions import ValidationError
+
+if TYPE_CHECKING:  # pragma: no cover
+    from collections.abc import Mapping
+    from pathlib import Path
+
+    from kreuzberg._types import ExtractionResult
+
+
+class ImageExtractor(Extractor):
+    SUPPORTED_MIME_TYPES: ClassVar[set[str]] = IMAGE_MIME_TYPES
+
+    IMAGE_MIME_TYPE_EXT_MAP: ClassVar[Mapping[str, str]] = {
+        "image/bmp": "bmp",
+        "image/x-bmp": "bmp",
+        "image/x-ms-bmp": "bmp",
+        "image/gif": "gif",
+        "image/jpeg": "jpg",
+        "image/pjpeg": "jpg",
+        "image/png": "png",
+        "image/tiff": "tiff",
+        "image/x-tiff": "tiff",
+        "image/jp2": "jp2",
+        "image/jpx": "jpx",
+        "image/jpm": "jpm",
+        "image/mj2": "mj2",
+        "image/webp": "webp",
+        "image/x-portable-anymap": "pnm",
+        "image/x-portable-bitmap": "pbm",
+        "image/x-portable-graymap": "pgm",
+        "image/x-portable-pixmap": "ppm",
+    }
+
+    async def extract_bytes_async(self, content: bytes) -> ExtractionResult:
+        extension = self._get_extension_from_mime_type(self.mime_type)
+        file_path, unlink = await create_temp_file(f".{extension}")
+        await AsyncPath(file_path).write_bytes(content)
+        try:
+            return await self.extract_path_async(file_path)
+        finally:
+            await unlink()
+
+    async def extract_path_async(self, path: Path) -> ExtractionResult:
+        if self.config.ocr_backend is None:
+            raise ValidationError("ocr_backend is None, cannot perform OCR")
+
+        return await get_ocr_backend(self.config.ocr_backend).process_file(path, **self.config.get_config_dict())
+
+    def extract_bytes_sync(self, content: bytes) -> ExtractionResult:
+        return anyio.run(self.extract_bytes_async, content)
+
+    def extract_path_sync(self, path: Path) -> ExtractionResult:
+        return anyio.run(self.extract_path_async, path)
+
+    def _get_extension_from_mime_type(self, mime_type: str) -> str:
+        if mime_type in self.IMAGE_MIME_TYPE_EXT_MAP:
+            return self.IMAGE_MIME_TYPE_EXT_MAP[mime_type]
+
+        for k, v in self.IMAGE_MIME_TYPE_EXT_MAP.items():
+            if k.startswith(mime_type):
+                return v
+
+        raise ValidationError("unsupported mimetype", context={"mime_type": mime_type})